import tweb

# Chat room storage - in production, you might want to use Redis or a database
chat_sessions = set()
chat_history = []

@tweb.websocket_handler("/chat")
class ChatWebSocketHandler(tweb.WebsocketHandler):
    def on_open(self, session):
        # Add session to the chat room
        chat_sessions.add(session)
        # Send chat history to the new user
        for message in chat_history[-20:]:  # Send last 20 messages
            session.send_text(message)
        # Notify all users about new connection
        broadcast_message(f"User joined the chat. ({len(chat_sessions)} users online)")

    def on_text_message(self, session, message):
        # Create a message with user info
        user_message = f"User {session.id[:8]}: {message}"
        # Store message in history
        chat_history.append(user_message)
        # Keep only last 100 messages
        if len(chat_history) > 100:
            chat_history.pop(0)
        # Broadcast message to all users
        broadcast_message(user_message)

    def on_close(self, session, reason):
        # Remove session from chat room
        chat_sessions.discard(session)
        # Notify all users about disconnection
        broadcast_message(f"User left the chat. ({len(chat_sessions)} users online)")

def broadcast_message(message):
    # Send message to all connected sessions
    for session in list(chat_sessions):
        if not session.is_closed:
            try:
                session.send_text(message)
            except:
                # Remove dead sessions
                chat_sessions.discard(session)